import bisect
from cProfile import label
from re import T
import warnings
from torchvision import transforms
from PIL import Image
import cv2
import numpy as np
import pandas as pd
import torch
import torch.utils.data as data

# from train import main

# def Split(image):
    
#     B,G,R = cv2.split(image)

#     B_equal = cv2.equalizeHist(B)

#     G_equal = cv2.equalizeHist(G)

#     R_equal = cv2.equalizeHist(R)


#     Result_0 = cv2.merge([B_equal,G_equal,R_equal])
#     Result_1 = cv2.hconcat((image,Result_0))

#     return Result_0;

# def horizontal_flip(image):
#     HF = transforms.RandomHorizontalFlip()
#     hf_image = HF(image)
#     return hf_image

# def vertical_flip(image):
#     VF = transforms.RandomVerticalFlip()
#     vf_image = VF(image)
#     return vf_image


# def random_rotation(image):
#     RR = transforms.RandomRotation(degrees=(0, 80))
#     rr_image = RR(image)
#     return rr_image


# class FaceDataset(data.Dataset):
#     # 初始化
#     def __init__(self, root,train=False):
#         super(FaceDataset, self).__init__()
#         self.root = root
#         self.train = train
#         df_path = pd.read_csv(root + '/dataset.csv', header=None, usecols=[0])
#         df_label = pd.read_csv(root + '/dataset.csv', header=None, usecols=[1])
#         # print("df_path", len(df_path))
#         self.path = np.array(df_path)[:, 0]
#         self.label = np.array(df_label)[:, 0]

#     # 读取某幅图片，item为索引号
#     def __getitem__(self, item):
#         # face = cv2.imread(self.root + '/' + self.path[item])
#         # face_gray = cv2.cvtColor(face, cv2.COLOR_BGR2GRAY)
#         # face_hist = cv2.equalizeHist(face_gray)
#         face = Image.open(self.root + '/' + self.path[item])
#         if self.train:
#             face = horizontal_flip(face)
#             face = vertical_flip(face)
#             face = random_rotation(face)
#         # 图像均衡化
#         face_hist = Split(face)
#         face_hist = np.array(face)
#         # """
#         face_normalized = face_hist.reshape(1, 100, 100) / 255.0
#         face_tensor = torch.from_numpy(face_normalized)
#         face_tensor = face_tensor.type('torch.FloatTensor')
#         label = self.label[item]
#         return face_tensor, int(label)

#     # 获取数据集样本个数
#     def __len__(self):
#         return self.path.shape[0]

# def get_alllabels_from_dataset(item):
#     return item[1]

class FaceDataset(data.Dataset):
    # 初始化
    def __init__(self, root):
        super(FaceDataset, self).__init__()
        self.root = root
        df_path = pd.read_csv(root + '/dataset.csv', header=None, usecols=[1])[1:]
        df_label = pd.read_csv(root + '/dataset.csv', header=None, usecols=[2])[1:]
        self.path = np.array(df_path)[:, 0]
        self.label = np.array(df_label)[:, 0]

    def __getitem__(self, item):
        face = cv2.imread(self.root + '/' + self.path[item])
        # 读取单通道灰度图
        face_gray = cv2.cvtColor(face, cv2.COLOR_BGR2GRAY)
        # 直方图均衡化
        face_hist = cv2.equalizeHist(face_gray)

        face_normalized = face_hist.reshape(1, 48, 48) / 255.0
        face_tensor = torch.from_numpy(face_normalized)
        face_tensor = face_tensor.type('torch.FloatTensor')
        label = self.label[item]
        return face_tensor, int(label)

    # 获取数据集样本个数
    def __len__(self):
        return self.path.shape[0]


if __name__ == '__main__':
    df_path = pd.read_csv('dataset/cnn_train/dataset.csv', header=None, usecols=[1])[1:]
    print(df_path)

